from datetime import date, datetime
from typing import List

import joblib

from pred_models import AGRUModel, GRUModel, MLPModel, get_gbdt_model
from trend_predict import TrainTestJob
from utils import get_all_stock_list_in_sector


def get_train_test_time(test_start: str, test_end: str) -> List[List[str]]:
    """
    将选定的测试时间范围按半年为单位分多个测试区块，并计算测试集起始时间
    :param test_start: "20230211"
    :param test_end: "20240610"
    :return: [["20190101", "20230101", "20230630"] ......]
    """
    start_date = datetime.strptime(test_start, "%Y%m%d").date()
    end_date = datetime.strptime(test_end, "%Y%m%d").date()
    pstart = date(year=start_date.year, month=1 if start_date.month <= 6 else 7, day=1)
    pend = date(
        year=start_date.year,
        month=6 if start_date.month <= 6 else 12,
        day=30 if start_date.month <= 6 else 31,
    )
    ptrain_start = date(year=pstart.year - 4, month=pstart.month, day=pstart.day)
    all_time_list = [[ptrain_start, pstart, pend]]
    while pend < end_date:
        if pstart.month <= 6:
            pstart = date(year=pstart.year, month=7, day=1)
            pend = date(year=pstart.year, month=12, day=31)
        else:
            xyear = pstart.year + 1
            pstart = date(year=xyear, month=1, day=1)
            pend = date(year=xyear, month=6, day=30)
        ptrain_start = date(year=pstart.year - 4, month=pstart.month, day=pstart.day)
        all_time_list.append([ptrain_start, pstart, pend])

    all_time_list = [
        [x.strftime("%Y%m%d"), y.strftime("%Y%m%d"), z.strftime("%Y%m%d")]
        for x, y, z in all_time_list
    ]
    return all_time_list


func_model_map = {0: get_gbdt_model, 1: MLPModel, 2: GRUModel, 3: AGRUModel}


def run_one_model(
    train_start: str,
    test_start: str,
    test_end: str,
    func_model_id: int,
    func_id: int,
    model_path: str = None,
):
    """
    跑一次实验
    :param train_start: 训练集起始时间，类似 ”20210101“
    :param test_start: 测试集始时间，类似 ”20220101“
    :param test_end: 测试集结束时间，类似 ”20220630“
    :param func_model_id: 使用的模型id
    :param func_id: 使用的数据处理函数的id
    :param model_path: 为模型加载数据
    :return:
    """
    print(train_start, test_start, test_end, func_model_id, func_id, model_path)
    code_market_list = get_all_stock_list_in_sector("上证A股")
    func_model = func_model_map.get(func_model_id)
    if func_model_id in [0, 1]:
        job_manager = TrainTestJob(
            model_define_func=func_model,
            code_market_list=code_market_list,
            start=train_start,
            train_split=test_start,
            end=test_end,
            period="1d",
            dividend_type="front_ratio",
            pred_window=10,
            history_days=1,
            history_window=30,
            func_id=func_id,
            dir_model_base=f"data_results8/{func_model.__name__}",
        )
    else:
        job_manager = TrainTestJob(
            model_define_func=func_model,
            code_market_list=code_market_list,
            start=train_start,
            train_split=test_start,
            end=test_end,
            period="1d",
            dividend_type="front_ratio",
            pred_window=10,
            history_days=30,
            history_window=1,
            func_id=func_id,
            dir_model_base=f"data_results8/{func_model.__name__}",
        )
    if not model_path:
        job_manager.train()
    job_manager.evaluate(is_test=True, model_file=model_path)
    # job_manager.evaluate(is_test=False, model_file=model_path)


def run_by_quanters(test_start: str, test_end: str, func_model_id: int, func_id: int):
    """
    把指定时间范围作为测试集，将测试集按照半年（自然年）划分区块作为测试集，并在测试集之前选取4年时间作为训练集。然后在指定模型，
    和指定的数据处理方式上，对模型进行训练和测试，并保存测试结果。
    :param test_start: 起始时间，”20210101“
    :param test_end: 结束时间，”20210101“
    :param func_model_id: 模型id
    :param func_id: 函数处理方式，和custom_dataset中定义相似
    :return:
    """
    all_time_list = get_train_test_time(test_start, test_end)  # 划分测试集区块
    all_time_list = list(all_time_list.__reversed__())
    pool = joblib.Parallel(n_jobs=4)
    pool(
        joblib.delayed(run_one_model)(x, y, z, func_model_id, func_id)
        for x, y, z in all_time_list
    )


def run_by_quanters_load_models(
    test_start: str, test_end: str, func_model_id: int, func_id: int
):
    """
    把指定时间范围作为测试集，将测试集按照半年（自然年）划分区块作为测试集，并在测试集之前选取4年时间作为训练集。然后在指定模型，
    和指定的数据处理方式上，对模型进行训练和测试，并保存测试结果。
    :param test_start: 起始时间，”20210101“
    :param test_end: 结束时间，”20210101“
    :param func_model_id: 模型id
    :param func_id: 函数处理方式，和custom_dataset中定义相似
    :return:
    """
    model_time_to_model_path = {
        1: {
            (
                "20210101",
                "20210630",
            ): "data_results5/MLPModel/20240805000053_969318/model.pkl",
            (
                "20210701",
                "20211231",
            ): "data_results5/MLPModel/20240805003833_bba58f/model.pkl",
            (
                "20220101",
                "20220630",
            ): "data_results5/MLPModel/20240805003830_62ac1e/model.pkl",
            (
                "20220701",
                "20221231",
            ): "data_results5/MLPModel/20240805003830_3208e3/model.pkl",
            (
                "20230101",
                "20230630",
            ): "data_results5/MLPModel/20240804231601_f98ff7/model.pkl",
            (
                "20230701",
                "20231231",
            ): "data_results5/MLPModel/20240805002654_7743aa/model.pkl",
            (
                "20240101",
                "20240630",
            ): "data_results5/MLPModel/20240805002654_999e1a/model.pkl",
        },
        2: {
            (
                "20210101",
                "20210630",
            ): "data_results5/GRUModel/20240804235337_6d6aab/model.pkl",
            (
                "20210701",
                "20211231",
            ): "data_results5/GRUModel/20240804232333_850aab/model.pkl",
            (
                "20220101",
                "20220630",
            ): "data_results5/GRUModel/20240804232333_f3fec7/model.pkl",
            (
                "20220701",
                "20221231",
            ): "data_results5/GRUModel/20240804232333_03abaa/model.pkl",
            (
                "20230101",
                "20230630",
            ): "data_results5/GRUModel/20240804221042_7dab47/model.pkl",
            (
                "20230701",
                "20231231",
            ): "data_results5/GRUModel/20240805002936_d544af/model.pkl",
            (
                "20240101",
                "20240630",
            ): "data_results5/GRUModel/20240805002936_5b4877/model.pkl",
        },
        3: {
            (
                "20210101",
                "20210630",
            ): "data_results5/AGRUModel/20240805000031_0e629f/model.pkl",
            (
                "20210701",
                "20211231",
            ): "data_results5/AGRUModel/20240804235748_ce6498/model.pkl",
            (
                "20220101",
                "20220630",
            ): "data_results5/AGRUModel/20240804235748_c33fd8/model.pkl",
            (
                "20220701",
                "20221231",
            ): "data_results5/AGRUModel/20240804235748_6bc6f2/model.pkl",
            (
                "20230101",
                "20230630",
            ): "data_results5/AGRUModel/20240804224232_798ee1/model.pkl",
            (
                "20230701",
                "20231231",
            ): "data_results5/AGRUModel/20240805003614_c56884/model.pkl",
            (
                "20240101",
                "20240630",
            ): "data_results5/AGRUModel/20240805003613_0f8336/model.pkl",
        },
    }
    all_time_list = get_train_test_time(test_start, test_end)  # 划分测试集区块
    all_time_list = list(all_time_list.__reversed__())
    pool = joblib.Parallel(n_jobs=1)
    pool(
        joblib.delayed(run_one_model)(
            x,
            y,
            z,
            func_model_id,
            func_id,
            model_time_to_model_path.get(func_model_id, {}).get((y, z)),
        )
        for x, y, z in all_time_list
    )


def draw_fig():
    file_df = r"data_model_xiaopan/20240719145836_58df03/test/pred_result.csv"
    TrainTestJob.view_result_data(file_df, "SA00.ZF", "20240101", "20240114")


if __name__ == "__main__":
    # run_by_quanters("20210101", "20240630", 2, 1)
    # run_by_quanters("20210101", "20240630", 3, 1)
    # run_by_quanters("20210101", "20240630", 1, 1)
    # run_by_quanters("20210101", "20240630", 0, 0)
    import sys

    st, ed, model, data_func = sys.argv[1:]
    run_by_quanters(st, ed, int(model), int(data_func))
